import torch
import torch.nn as nn

from enum import Enum
from typing import Dict
from torchvision import models
from torchvision.models.segmentation import DeepLabV3_ResNet101_Weights

from .deeplabV3Plus import DeeplabV3Plus


class ModelType(Enum):
    DEEPLAB_V3_PLUS = 'deeplab_v3_plus'
    DEEPLAB_V3 = 'deeplab_v3'


def get_model(cfg: Dict[str, dict]):
    model_type = cfg.get('model')
    model = None

    if model_type == ModelType.DEEPLAB_V3_PLUS:
        model = DeeplabV3Plus(cfg)
    elif model_type == ModelType.DEEPLAB_V3:
        model = models.segmentation.deeplabv3_resnet101(weights=DeepLabV3_ResNet101_Weights.DEFAULT)
        model.classifier[-1] = nn.Conv2d(256, cfg.get('num_classes'), kernel_size=1)
    
    if torch.cuda.device_count() > 1:
        model = nn.DataParallel(model)
    return model
